import pandas as pd
import os
import glob


# 方法1：使用glob获取所有csv文件
def load_all_csv_files(directory_path):
    # 获取目录下所有csv文件路径
    csv_files = glob.glob(os.path.join(directory_path, "*.csv"))

    # 创建一个空字典来存储所有DataFrame
    dataframes = {}

    for file_path in csv_files:
        # 提取文件名（不含扩展名）作为key
        file_name = os.path.basename(file_path).split('.')[0]

        try:
            # 读取CSV文件
            df = pd.read_csv(file_path)
            dataframes[file_name] = df
            print(f"成功加载: {file_name}")
        except Exception as e:
            print(f"加载失败 {file_name}: {e}")

    return dataframes


# 使用示例
# directory = "./data"  # 替换为你的目录路径
# all_data = load_all_csv_files(directory)



def load_files_by_pattern(file_pattern):
    """
    根据文件模式加载多个文件

    参数:
    file_pattern: 文件模式，如 "data/*.csv" 或 "data/*.xlsx"
    """
    file_paths = glob.glob(file_pattern)
    dataframes = {}

    for file_path in file_paths:
        file_name = file_path.split('/')[-1].split('.')[0]

        try:
            if file_path.endswith('.csv'):
                df = pd.read_csv(file_path)
            elif file_path.endswith('.xlsx') or file_path.endswith('.xls'):
                df = pd.read_excel(file_path)
            elif file_path.endswith('.json'):
                df = pd.read_json(file_path)
            else:
                print(f"不支持的文件格式: {file_path}")
                continue

            dataframes[file_name] = df
            print(f"成功加载: {file_name}")

        except Exception as e:
            print(f"加载失败 {file_path}: {e}")

    return dataframes


# 使用示例
# csv_data = load_files_by_pattern("data/*.csv")
# excel_data = load_files_by_pattern("data/*.xlsx")


def merge_csv_files(directory_path):
    """
    将目录下所有CSV文件合并为一个DataFrame
    """
    csv_files = glob.glob(os.path.join(directory_path, "*.csv"))

    dfs = []
    for file_path in csv_files:
        try:
            df = pd.read_csv(file_path)
            # 添加文件名作为新列，便于识别数据来源
            df['source_file'] = os.path.basename(file_path)
            dfs.append(df)
            print(f"成功加载: {os.path.basename(file_path)}")
        except Exception as e:
            print(f"加载失败 {file_path}: {e}")

    if dfs:
        # 合并所有DataFrame
        merged_df = pd.concat(dfs, ignore_index=True)
        return merged_df
    else:
        return pd.DataFrame()


# 使用示例
# combined_data = merge_csv_files("./data")
# print(f"合并后的数据形状: {combined_data.shape}")\

# import pandas as pd
# import os


def load_files_os(directory_path, file_extension='.csv'):
    """
    使用os.listdir()加载目录下的文件
    """
    dataframes = {}

    for filename in os.listdir(directory_path):
        if filename.endswith(file_extension):
            file_path = os.path.join(directory_path, filename)
            file_key = filename.replace(file_extension, '')

            try:
                if file_extension == '.csv':
                    df = pd.read_csv(file_path)
                elif file_extension in ['.xlsx', '.xls']:
                    df = pd.read_excel(file_path)
                else:
                    continue

                dataframes[file_key] = df
                print(f"成功加载: {filename}")

            except Exception as e:
                print(f"加载失败 {filename}: {e}")

    return dataframes


# 使用示例
# data = load_files_os("./data", '.csv')

import pandas as pd
import glob
import os


def merge_xlsx_files(directory_path):
    """
    合并目录下所有xlsx文件为一个DataFrame
    """
    # 获取所有xlsx文件
    xlsx_files = glob.glob(os.path.join(directory_path, "*.xlsx"))

    dfs = []
    for file_path in xlsx_files:
        try:
            # 读取整个Excel文件（如果有多个sheet，会读取第一个）
            df = pd.read_excel(file_path)
            # 添加文件名作为标识列
            df['source_file'] = os.path.basename(file_path)
            dfs.append(df)
            print(f"成功加载: {os.path.basename(file_path)}")
        except Exception as e:
            print(f"加载失败 {file_path}: {e}")

    if dfs:
        merged_df = pd.concat(dfs, ignore_index=True)
        print(f"合并完成，总行数: {len(merged_df)}")
        return merged_df
    else:
        return pd.DataFrame()

# combined_data = merge_xlsx_files("../words")
# print(f"合并后的数据形状: {combined_data.shape}")